"""
导入拉勾招聘数据命令
"""
import pandas as pd
from django.core.management.base import BaseCommand
from django.conf import settings
from apps.positions.models import Position
from utils.data_clean import clean_salary, clean_html, clean_text_field, validate_salary


class Command(BaseCommand):
    help = '导入拉勾招聘CSV数据到数据库'

    def handle(self, *args, **options):
        self.stdout.write('='* 50)
        self.stdout.write('开始导入拉勾招聘数据...')
        self.stdout.write('='* 50)

        # CSV文件路径
        csv_path = settings.DATA_DIR / settings.LAGOU_CSV_FILE
        
        if not csv_path.exists():
            self.stdout.write(self.style.ERROR(f'CSV文件不存在: {csv_path}'))
            return

        try:
            # 读取CSV
            self.stdout.write(f'正在读取CSV文件: {csv_path}')
            df = pd.read_csv(
                csv_path,
                encoding='utf-8',
                dtype={'id': int, 'positionId': int, 'companyId': int},
            )
            self.stdout.write(self.style.SUCCESS(f'成功读取 {len(df)} 条数据'))
        except Exception as e:
            self.stdout.write(self.style.ERROR(f'读取CSV失败: {e}'))
            return

        # 数据清洗
        self.stdout.write('开始数据清洗...')
        df = self.clean_data(df)
        self.stdout.write(self.style.SUCCESS('数据清洗完成'))

        # 导入数据库
        self.stdout.write('开始导入数据库...')
        positions = []
        success_count = 0
        error_count = 0
        
        for idx, row in df.iterrows():
            try:
                position = Position(
                    original_id=int(row['id']) if pd.notna(row['id']) else None,
                    category=row.get('category'),
                    big_category=row.get('bigcategory'),
                    position_id=int(row['positionId']) if pd.notna(row.get('positionId')) else None,
                    position_name=row['positionName'],
                    company_id=int(row['companyId']) if pd.notna(row.get('companyId')) else None,
                    company_full_name=row.get('companyFullName'),
                    company_size=row.get('companySize'),
                    industry_field=row.get('industryField'),
                    finance_stage=row.get('financeStage'),
                    position_type=row.get('positionType'),
                    city=row.get('city'),
                    district=row.get('district'),
                    salary=row.get('salary'),
                    salary_min=row.get('salary_min'),
                    salary_max=row.get('salary_max'),
                    salary_avg=row.get('salary_avg'),
                    work_year=row.get('workYear'),
                    job_nature=row.get('jobNature'),
                    education=row.get('education'),
                    position_detail=row.get('position_detail_clean'),  # 使用清洗后的
                    position_advantage=row.get('positionAdvantage'),
                )
                positions.append(position)
                success_count += 1

                # 批量插入（每1000条）
                if len(positions) >= 1000:
                    Position.objects.bulk_create(positions, ignore_conflicts=True)
                    self.stdout.write(f'已导入 {success_count} 条数据')
                    positions = []

            except Exception as e:
                error_count += 1
                if error_count <= 5:  # 只显示前5个错误
                    self.stdout.write(self.style.WARNING(f'第 {idx + 1} 行导入失败: {e}'))

        # 插入剩余数据
        if positions:
            Position.objects.bulk_create(positions, ignore_conflicts=True)

        self.stdout.write('='* 50)
        self.stdout.write(self.style.SUCCESS(
            f'数据导入完成！成功: {success_count}, 失败: {error_count}'
        ))
        self.stdout.write('='* 50)

    def clean_data(self, df):
        """数据清洗"""
        # 1. 解析薪资
        df['salary_min'] = None
        df['salary_max'] = None
        df['salary_avg'] = None
        
        for idx, row in df.iterrows():
            salary_str = row.get('salary')
            if pd.notna(salary_str):
                salary_min, salary_max, salary_avg, _ = clean_salary(salary_str)
                if validate_salary(salary_min, salary_max):
                    df.at[idx, 'salary_min'] = salary_min
                    df.at[idx, 'salary_max'] = salary_max
                    df.at[idx, 'salary_avg'] = salary_avg
        
        # 2. 清理HTML标签（职位详情）
        if 'positionDetail' in df.columns:
            df['position_detail_clean'] = df['positionDetail'].apply(clean_html)
        
        # 3. 清理文本字段
        text_fields = ['positionName', 'companyFullName', 'city', 'district', 'education', 'workYear']
        for field in text_fields:
            if field in df.columns:
                df[field] = df[field].apply(clean_text_field)
        
        # 4. 去除关键字段为空的记录
        df = df.dropna(subset=['positionName', 'city'])
        
        return df

